import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import odeint

# ====================== 修复1：中文显示问题 ======================
# 设置支持中文的字体（添加更多备选字体）
plt.rcParams['font.sans-serif'] = ['Microsoft YaHei', 'SimHei', 'KaiTi', 'FangSong', 'Arial Unicode MS']
plt.rcParams['axes.unicode_minus'] = False

# 设置随机种子以确保结果可重现
np.random.seed(42)

# 无人机模型参数
m = 1.0      # 质量 (kg)
g = 9.8      # 重力加速度 (m/s^2)
c = 0.2      # 空气阻力系数

# 非线性ADRC函数
def fal(e, alpha, delta):
    """非线性函数，用于ESO和控制器"""
    if abs(e) <= delta:
        return e / (delta**(1-alpha))
    else:
        return abs(e)**alpha * np.sign(e)

# 扩张状态观测器 (ESO) 类
class ESO:
    def __init__(self, n, beta, alpha, delta, dt):
        self.n = n  # 系统阶数
        self.beta = beta  # 观测器增益
        self.alpha = alpha  # 非线性参数
        self.delta = delta  # 线性区间宽度
        self.dt = dt  # 采样时间
        self.z = np.zeros(n+1)  # 状态估计 [z1, z2, ..., z_{n+1}]
        
    def update(self, y, u):
        # 计算观测误差
        e = y - self.z[0]
        
        # 更新状态估计
        for i in range(self.n):
            self.z[i] += self.dt * (self.z[i+1] + self.beta[i] * fal(e, self.alpha[i], self.delta))
        
        # 更新扩张状态（总扰动）
        self.z[self.n] += self.dt * self.beta[self.n] * fal(e, self.alpha[self.n], self.delta)
        
        return self.z

# 非线性状态误差反馈控制律
def nlsef(e, k, alpha, delta):
    """非线性状态误差反馈"""
    u = 0
    for i in range(len(e)):
        u += k[i] * fal(e[i], alpha[i], delta)
    return u

# 无人机动力学模型
def drone_dynamics(x, t, u, disturbance):
    h, v = x  # 高度, 速度
    
    # 系统动力学方程
    dhdt = v
    dvdt = u/m - g - (c/m)*v + disturbance(t)
    
    return [dhdt, dvdt]

# 外部扰动函数
def disturbance(t):
    # 时变扰动
    if t > 15 and t < 20:
        return 2.0 * np.sin(0.5*t)  # 正弦扰动
    elif t > 25:
        return -1.5  # 阶跃扰动
    return 0.0

# 仿真参数
T = 30.0  # 总仿真时间
dt = 0.01  # 采样时间
N = int(T/dt)  # 仿真步数
t = np.linspace(0, T, N)  # 时间向量

# 初始化数组存储结果
h_ref = np.zeros(N)  # 参考高度
h_actual = np.zeros(N)  # 实际高度
v_actual = np.zeros(N)  # 实际速度
u_total = np.zeros(N)  # 总控制量
u_inner = np.zeros(N)  # 内环控制量
disturbances = np.zeros(N)  # 扰动
h_estimate = np.zeros(N)  # 高度估计
v_estimate = np.zeros(N)  # 速度估计
disturbance_estimate = np.zeros(N)  # 扰动估计

# 参考信号 - 多段高度指令
for i in range(N):
    if t[i] < 5:
        h_ref[i] = 5.0
    elif t[i] < 15:
        h_ref[i] = 15.0
    elif t[i] < 25:
        h_ref[i] = 8.0
    else:
        h_ref[i] = 12.0

# 初始化ADRC控制器 - 优化后的参数

# 外环ESO (高度环)
eso_outer = ESO(
    n=2,  # 系统阶数
    beta=[150, 500, 2000],  # 增加观测器增益以提高估计速度
    alpha=[0.5, 0.75, 1.25],  # 非线性参数
    delta=0.05,  # 减小线性区间宽度以提高精度
    dt=dt
)

# 内环ESO (速度环)
eso_inner = ESO(
    n=1,  # 系统阶数
    beta=[300, 1000],  # 增加观测器增益
    alpha=[0.5, 1.25],  # 非线性参数
    delta=0.05,  # 减小线性区间宽度
    dt=dt
)

# 初始状态
x = [0.0, 0.0]  # [高度, 速度]

# 仿真循环
for i in range(1, N):
    # 记录当前扰动
    disturbances[i] = disturbance(t[i])
    
    # 外环ADRC (高度控制)
    z_outer = eso_outer.update(x[0], u_inner[i-1])
    e1 = h_ref[i] - z_outer[0]  # 高度误差
    e2 = 0 - z_outer[1]  # 高度微分误差
    v_ref = nlsef([e1, e2], [0.8, 1.2], [0.5, 0.75], 0.05) - z_outer[2]/eso_outer.beta[0]
    
    # 记录估计值
    h_estimate[i] = z_outer[0]
    disturbance_estimate[i] = z_outer[2]
    
    # 内环ADRC (速度控制)
    z_inner = eso_inner.update(x[1], u_total[i-1])
    e_v = v_ref - z_inner[0]  # 速度误差
    u0 = nlsef([e_v], [8.0], [0.75], 0.05)  # 增加控制增益
    u_inner[i] = u0 - z_inner[1]/eso_inner.beta[0]
    
    # 记录速度估计
    v_estimate[i] = z_inner[0]
    
    # 饱和限制
    u_inner[i] = np.clip(u_inner[i], -25, 25)  # 增加控制限幅
    u_total[i] = u_inner[i] * m  # 转换为力
    
    # 更新系统状态
    t_range = [t[i-1], t[i]]
    x = odeint(drone_dynamics, x, t_range, args=(u_total[i], disturbance))[1]
    
    # 记录结果
    h_actual[i] = x[0]
    v_actual[i] = x[1]

# 绘制结果
plt.figure(figsize=(15, 12))

# 高度响应
plt.subplot(4, 2, 1)
plt.plot(t, h_actual, 'r-', label='实际高度')
plt.plot(t, h_ref, 'b--', label='参考高度')
plt.plot(t, h_estimate, 'g:', label='估计高度', alpha=0.7)
plt.ylabel('高度 (m)')
plt.legend(loc='upper right')
plt.grid(True)
plt.title('高度响应')

# 速度响应
plt.subplot(4, 2, 2)
plt.plot(t, v_actual, 'm-', label='实际速度')
plt.plot(t, v_estimate, 'c:', label='估计速度', alpha=0.7)
plt.ylabel('速度 (m/s)')
plt.legend(loc='upper right')
plt.grid(True)
plt.title('速度响应')

# 控制输入
plt.subplot(4, 2, 3)
plt.plot(t, u_total, 'g-')
plt.ylabel('控制力 (N)')
plt.grid(True)
plt.title('控制输入')

# 扰动与估计
plt.subplot(4, 2, 4)
plt.plot(t, disturbances, 'k-', label='实际扰动')
plt.plot(t, disturbance_estimate, 'r:', label='估计扰动', alpha=0.7)
plt.xlabel('时间 (s)')
plt.ylabel('扰动')
plt.legend(loc='upper right')
plt.grid(True)
plt.title('外部扰动与估计')

# 高度误差
plt.subplot(4, 2, 5)
height_error = h_ref - h_actual
plt.plot(t, height_error, 'r-')
plt.xlabel('时间 (s)')
plt.ylabel('高度误差 (m)')
plt.grid(True)
plt.title('高度跟踪误差')

# 控制输入细节
plt.subplot(4, 2, 6)
plt.plot(t, u_total, 'g-')
plt.xlabel('时间 (s)')
plt.ylabel('控制力 (N)')
plt.grid(True)
plt.ylim(-30, 30)
plt.title('控制输入 (细节)')

# 相位图
plt.subplot(4, 2, 7)
plt.plot(h_actual, v_actual, 'b-')
plt.xlabel('高度 (m)')
plt.ylabel('速度 (m/s)')
plt.grid(True)
plt.title('相位图: 高度 vs 速度')

# 扰动估计误差
plt.subplot(4, 2, 8)
disturbance_error = disturbances - disturbance_estimate
plt.plot(t, disturbance_error, 'r-')
plt.xlabel('时间 (s)')
plt.ylabel('扰动估计误差')
plt.grid(True)
plt.title('扰动估计误差')

plt.tight_layout()
plt.savefig('adrc_performance.png', dpi=300)
plt.show()

# 计算性能指标
# 找到每个阶段的起始和结束索引
stage1_end = np.where(t >= 5)[0][0]
stage2_end = np.where(t >= 15)[0][0]
stage3_end = np.where(t >= 25)[0][0]

# 计算每个阶段的ISE
ise1 = np.sum((h_ref[:stage1_end] - h_actual[:stage1_end])**2) * dt
ise2 = np.sum((h_ref[stage1_end:stage2_end] - h_actual[stage1_end:stage2_end])**2) * dt
ise3 = np.sum((h_ref[stage2_end:stage3_end] - h_actual[stage2_end:stage3_end])**2) * dt
ise4 = np.sum((h_ref[stage3_end:] - h_actual[stage3_end:])**2) * dt

# 计算IAE (积分绝对误差)
iae1 = np.sum(np.abs(h_ref[:stage1_end] - h_actual[:stage1_end])) * dt
iae2 = np.sum(np.abs(h_ref[stage1_end:stage2_end] - h_actual[stage1_end:stage2_end])) * dt
iae3 = np.sum(np.abs(h_ref[stage2_end:stage3_end] - h_actual[stage2_end:stage3_end])) * dt
iae4 = np.sum(np.abs(h_ref[stage3_end:] - h_actual[stage3_end:])) * dt

# 计算ITAE (时间加权积分绝对误差)
itae1 = np.sum(t[:stage1_end] * np.abs(h_ref[:stage1_end] - h_actual[:stage1_end])) * dt
itae2 = np.sum(t[stage1_end:stage2_end] * np.abs(h_ref[stage1_end:stage2_end] - h_actual[stage1_end:stage2_end])) * dt
itae3 = np.sum(t[stage2_end:stage3_end] * np.abs(h_ref[stage2_end:stage3_end] - h_actual[stage2_end:stage3_end])) * dt
itae4 = np.sum(t[stage3_end:] * np.abs(h_ref[stage3_end:] - h_actual[stage3_end:])) * dt

print("性能指标:")
print(f"阶段1 (0-5s): ISE={ise1:.4f}, IAE={iae1:.4f}, ITAE={itae1:.4f}")
print(f"阶段2 (5-15s): ISE={ise2:.4f}, IAE={iae2:.4f}, ITAE={itae2:.4f}")
print(f"阶段3 (15-25s): ISE={ise3:.4f}, IAE={iae3:.4f}, ITAE={itae3:.4f}")
print(f"阶段4 (25-30s): ISE={ise4:.4f}, IAE={iae4:.4f}, ITAE={itae4:.4f}")
print(f"总指标: ISE={ise1+ise2+ise3+ise4:.4f}, IAE={iae1+iae2+iae3+iae4:.4f}, ITAE={itae1+itae2+itae3+itae4:.4f}")

# 计算扰动估计性能
disturbance_rmse = np.sqrt(np.mean((disturbances - disturbance_estimate)**2))
print(f"扰动估计RMSE: {disturbance_rmse:.4f}")
